import sys
import os
import numpy as np
import torch

from PySide6 import QtWidgets, QtCore
import pyqtgraph.opengl as gl
from dataclasses import dataclass

@dataclass
class Point3D:
    x: float
    y: float
    z: float

@dataclass
class RadarPoint(Point3D):
    strength: float
    velocity: float
    cohen_factor: float
    cohen_factor2: float = 0.0

def predict_cohenfactor2(points, model, device='cpu', batch_size=65536):
    if not points:
        return []
    feats = np.array([[p.x, p.y, p.z, p.strength] for p in points], dtype=np.float32)
    all_preds = []
    with torch.no_grad():
        for start in range(0, len(points), batch_size):
            batch = torch.from_numpy(feats[start:start+batch_size]).to(device)
            preds = model(batch).squeeze(-1).cpu().numpy()
            all_preds.append(preds)
    cohen2 = np.concatenate(all_preds)
    return cohen2

def add_cohenfactor2_to_points(points, model, device='cpu'):
    preds = predict_cohenfactor2(points, model, device=device)
    for p, c2 in zip(points, preds):
        p.cohen_factor2 = float(c2)
    return points

class SingleFramePointCloudViewer(QtWidgets.QWidget):
    def __init__(self, model, reader, file_list, device='cpu'):
        super().__init__()
        self.model = model
        self.reader = reader
        self.file_list = file_list
        self.device = device

        layout = QtWidgets.QVBoxLayout(self)

        self.view = gl.GLViewWidget()
        self.view.setCameraPosition(distance=10)
        layout.addWidget(self.view, stretch=10)

        self.frame_slider = QtWidgets.QSlider(QtCore.Qt.Horizontal)
        self.frame_slider.setMinimum(0)
        self.frame_slider.setMaximum(len(self.file_list)-1)
        self.frame_slider.setValue(0)
        layout.addWidget(QtWidgets.QLabel("帧号"))
        layout.addWidget(self.frame_slider)

        # 新增逻辑选项
        self.logic_combo = QtWidgets.QComboBox()
        self.logic_combo.addItems(["与(AND)", "或(OR)"])
        layout.addWidget(QtWidgets.QLabel("筛选逻辑"))
        layout.addWidget(self.logic_combo)

        sliders_layout = QtWidgets.QHBoxLayout()
        self.s_strength = self.make_slider("Strength min", sliders_layout)
        self.s_cohen = self.make_slider("Cohen Factor min", sliders_layout)
        self.s_cohen2 = self.make_slider("CohenFactor2 min", sliders_layout)
        layout.addLayout(sliders_layout)

        self.scatter = gl.GLScatterPlotItem()
        self.view.addItem(self.scatter)

        self.cur_points = []
        self.cur_xyz = self.cur_strength = self.cur_cohen = self.cur_cohen2 = self.cur_z = None

        self.frame_slider.valueChanged.connect(self.load_frame)
        self.logic_combo.currentIndexChanged.connect(self.update_points)
        self.s_strength.valueChanged.connect(self.update_points)
        self.s_cohen.valueChanged.connect(self.update_points)
        self.s_cohen2.valueChanged.connect(self.update_points)

        self.load_frame()

    def make_slider(self, name, layout):
        box = QtWidgets.QVBoxLayout()
        label = QtWidgets.QLabel(name)
        slider = QtWidgets.QSlider(QtCore.Qt.Vertical)
        slider.setMinimum(0)
        slider.setMaximum(100)
        slider.setValue(0)
        box.addWidget(label)
        box.addWidget(slider)
        layout.addLayout(box)
        return slider

    def load_frame(self):
        idx = self.frame_slider.value()
        file_path = self.file_list[idx]
        points = self.reader._read_single_file(file_path)
        if points:
            add_cohenfactor2_to_points(points, self.model, device=self.device)
            self.cur_points = points
            self.cur_xyz = np.array([[p.x, p.y, p.z] for p in points], dtype=np.float32)
            self.cur_strength = np.array([p.strength for p in points], dtype=np.float32)
            self.cur_cohen = np.array([p.cohen_factor for p in points], dtype=np.float32)
            self.cur_cohen2 = np.array([p.cohen_factor2 for p in points], dtype=np.float32)
            self.cur_z = self.cur_xyz[:, 2]
            for s in (self.s_strength, self.s_cohen, self.s_cohen2):
                s.setValue(0)
        else:
            self.cur_xyz = self.cur_strength = self.cur_cohen = self.cur_cohen2 = self.cur_z = None
        self.update_points()

    def update_points(self):
        idx = self.frame_slider.value()
        file_path = self.file_list[idx]
        fname = os.path.basename(file_path)

        
        ms = int(fname.split('_')[-1][:-4])
        total_seconds = ms // 1000
        milliseconds = ms % 1000
        hours = total_seconds // 3600
        minutes = (total_seconds % 3600) // 60
        seconds = total_seconds % 60
        if self.cur_xyz is None or len(self.cur_xyz) == 0:
            self.scatter.setData(pos=np.zeros((1,3)), size=6, color=(1,0,0,0.1))

            
            self.setWindowTitle(f"帧 {idx} | {fname} | {hours:02}:{minutes:02}:{seconds:02}.{milliseconds:03} | 无点")
            return

        s_th = np.percentile(self.cur_strength, self.s_strength.value())
        c_th = np.percentile(self.cur_cohen, self.s_cohen.value())
        c2_th = np.percentile(self.cur_cohen2, self.s_cohen2.value())
        m1 = self.cur_strength >= s_th
        m2 = self.cur_cohen   >= c_th
        m3 = self.cur_cohen2  >= c2_th

        # 逻辑合并
        if self.logic_combo.currentIndex() == 0:   # "与"
            mask = m1 & m2 & m3
        else:  # "或"
            mask = m1 | m2 | m3

        pts = self.cur_xyz[mask]
        z = self.cur_z[mask]
        if pts.shape[0] == 0:
            pts = np.zeros((1, 3), dtype=np.float32)
            z = np.zeros((1,), dtype=np.float32)

        z_norm = (z - z.min()) / (np.ptp(z) + 1e-6)
        color = np.empty((pts.shape[0], 4), dtype=float)
        color[:, 0] = z_norm
        color[:, 1] = 1 - z_norm
        color[:, 2] = 0
        color[:, 3] = 0.8

        self.scatter.setData(pos=pts, size=4, color=color)
        self.setWindowTitle(f"帧 {idx} | {fname} | {hours:02}:{minutes:02}:{seconds:02}.{milliseconds:03} | 点数: {pts.shape[0]}")

if __name__ == '__main__':
    model_path = "cohen_fix.pt"
    device = 'cuda'
    model = torch.jit.load(model_path, map_location=device)

    sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
    from configs.config import load_config
    from data_process.radar_reader import RadarReader

    cfg = load_config("./configs/data_process.py")
    cfg['data_batch_size'] = 100
    reader = RadarReader(cfg)

    data_root = list(cfg.get('data_root_list').values())[0]
    folder_path = os.path.join(data_root, 'SARRadar')
    files = reader._get_file_list(folder_path)
    files.sort(key=lambda x: int(x.split('_')[-1].split('.')[0]))

    app = QtWidgets.QApplication(sys.argv)
    win = SingleFramePointCloudViewer(model, reader, files, device=device)
    win.show()
    sys.exit(app.exec())
